Learning Optimal Features via Partial Invariance

Learning models that are robust to distribution shifts is a key concern in the context of their real-life applicability. Invariant Risk Minimization (IRM) popular framework aims learn from multiple environments. The success IRM requires an important assumption: underlying causal mechanisms/features remain invariant across When not satisfied, we show can over-constrain predictor and remedy this, propose relaxation via partial invariance. In this work, theoretically highlight sub-optimality then demonstrate how learning partition training domains help improve models. Several experiments, conducted both linear settings as well with deep neural networks on tasks over language image data, allow us verify our conclusions.